
import os
from zhipuai import ZhipuAI
from dotenv import load_dotenv, find_dotenv

_ = load_dotenv(find_dotenv())
api_key = os.environ.get('ZHIPUAI_API_KEY')
if api_key is None:
    raise ValueError("API Key is not set in the .env file")
client = ZhipuAI(api_key=api_key)

def get_completion(prompt, model="glm-4-plus", temperature=0.01):
    messages = [
        {"role": "user", "content": "我昨天给小明买了三本书"},
        {"role": "assistant", "content": "书是不错的礼物哦。"},
        {"role": "user", "content": prompt}
    ]
    response = client.chat.completions.create(
        model=model,
        messages=messages,
        temperature=temperature  
    )
    return response.choices[0].message.content
  

prompt = "我昨天给谁买书了？"
res = get_completion(prompt, temperature=0.7)  
print(res)